"""
数据库管理模块
支持PostgreSQL和Redis，处理并发访问和数据存储
"""
import os
import json
import logging
from typing import List, Dict, Set, Optional, Tuple
from datetime import datetime
import hashlib
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, Boolean, Index, UniqueConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool
from sqlalchemy.exc import IntegrityError, OperationalError
import redis
from contextlib import contextmanager
import threading


Base = declarative_base()


class Topic(Base):
    """话题表"""
    __tablename__ = 'topics'
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    topic_text = Column(Text, nullable=False, index=True)
    source_keyword = Column(String(255), nullable=False, index=True)
    search_round = Column(Integer, nullable=False, default=1)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
    
    # 创建唯一索引，避免重复话题
    __table_args__ = (
        UniqueConstraint('topic_text', 'source_keyword', name='uq_topic_source'),
        Index('idx_topic_text', 'topic_text'),
        Index('idx_source_keyword', 'source_keyword'),
        Index('idx_topic_search_round', 'search_round'),
    )


class Keyword(Base):
    """关键词表"""
    __tablename__ = 'keywords'
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    keyword_text = Column(String(500), nullable=False, index=True)
    source_topic = Column(Text, nullable=True, index=True)
    search_source = Column(String(100), nullable=True, index=True)  # 搜索引擎来源
    search_round = Column(Integer, nullable=False, default=1)
    is_processed = Column(Boolean, default=False, index=True)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
    
    # 创建唯一索引，避免重复关键词
    __table_args__ = (
        UniqueConstraint('keyword_text', name='uq_keyword_text'),
        Index('idx_keyword_text', 'keyword_text'),
        Index('idx_source_topic', 'source_topic'),
        Index('idx_search_source', 'search_source'),
        Index('idx_keyword_search_round', 'search_round'),
        Index('idx_is_processed', 'is_processed'),
    )


class SearchHistory(Base):
    """搜索历史表"""
    __tablename__ = 'search_history'
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    search_query = Column(String(500), nullable=False, index=True)
    search_round = Column(Integer, nullable=False, default=1)
    result_count = Column(Integer, default=0)
    search_time = Column(DateTime, default=datetime.now)
    
    __table_args__ = (
        Index('idx_search_query', 'search_query'),
        Index('idx_history_search_round', 'search_round'),
        Index('idx_search_time', 'search_time'),
    )


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, config: Dict):
        """
        初始化数据库管理器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 数据库配置
        db_config = config.get('database', {})
        self.postgres_config = db_config.get('postgres', {})
        self.redis_config = db_config.get('redis', {})
        
        # 初始化PostgreSQL连接
        self._init_postgres()
        
        # 初始化Redis连接
        self._init_redis()
        
        # 线程本地存储
        self._local = threading.local()
    
    def _init_postgres(self):
        """初始化PostgreSQL连接"""
        try:
            # 构建连接字符串
            host = self.postgres_config.get('host', 'localhost')
            port = self.postgres_config.get('port', 5432)
            database = self.postgres_config.get('database', 'keyword_mining')
            username = self.postgres_config.get('username', 'postgres')
            password = self.postgres_config.get('password', '')
            
            connection_string = f"postgresql://{username}:{password}@{host}:{port}/{database}"
            
            # 创建引擎，配置连接池
            self.engine = create_engine(
                connection_string,
                poolclass=QueuePool,
                pool_size=self.postgres_config.get('pool_size', 10),
                max_overflow=self.postgres_config.get('max_overflow', 20),
                pool_pre_ping=True,
                pool_recycle=3600,
                echo=self.postgres_config.get('echo', False)
            )
            
            # 创建会话工厂
            self.SessionLocal = sessionmaker(
                bind=self.engine,
                autocommit=False,
                autoflush=False
            )
            
            # 创建表
            Base.metadata.create_all(bind=self.engine)
            
            self.logger.info("PostgreSQL连接初始化成功")
            
        except Exception as e:
            self.logger.error(f"PostgreSQL连接初始化失败: {e}")
            raise
    
    def _init_redis(self):
        """初始化Redis连接"""
        try:
            # 创建Redis连接池
            self.redis_pool = redis.ConnectionPool(
                host=self.redis_config.get('host', 'localhost'),
                port=self.redis_config.get('port', 6379),
                password=self.redis_config.get('password', None),
                db=self.redis_config.get('db', 0),
                max_connections=self.redis_config.get('max_connections', 20),
                decode_responses=True
            )
            
            self.redis_client = redis.Redis(connection_pool=self.redis_pool)
            
            # 测试连接
            self.redis_client.ping()
            
            self.logger.info("Redis连接初始化成功")
            
        except Exception as e:
            self.logger.error(f"Redis连接初始化失败: {e}")
            raise
    
    @contextmanager
    def get_session(self) -> Session:
        """
        获取数据库会话的上下文管理器
        
        Yields:
            SQLAlchemy会话对象
        """
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            self.logger.error(f"数据库操作失败: {e}")
            raise
        finally:
            session.close()
    
    def add_topic(self, topic_text: str, source_keyword: str, search_round: int) -> bool:
        """
        添加话题到数据库
        
        Args:
            topic_text: 话题文本
            source_keyword: 源关键词
            search_round: 搜索轮数
        
        Returns:
            是否添加成功
        """
        try:
            with self.get_session() as session:
                # 检查是否已存在
                existing = session.query(Topic).filter_by(
                    topic_text=topic_text,
                    source_keyword=source_keyword
                ).first()
                
                if existing:
                    return False
                
                # 创建新话题
                topic = Topic(
                    topic_text=topic_text,
                    source_keyword=source_keyword,
                    search_round=search_round
                )
                
                session.add(topic)
                session.commit()
                
                self.logger.debug(f"成功添加话题: {topic_text}")
                return True
                
        except IntegrityError:
            # 并发情况下可能出现的重复键错误
            return False
        except Exception as e:
            self.logger.error(f"添加话题失败: {e}")
            return False
    
    def add_keyword(self, keyword_text: str, source_topic: str = None, search_source: str = None, search_round: int = 1) -> bool:
        """
        添加关键词到数据库
        
        Args:
            keyword_text: 关键词文本
            source_topic: 源话题
            search_source: 搜索引擎来源
            search_round: 搜索轮数
        
        Returns:
            是否添加成功
        """
        try:
            with self.get_session() as session:
                # 检查是否已存在
                existing = session.query(Keyword).filter_by(keyword_text=keyword_text).first()
                
                if existing:
                    return False
                
                # 创建新关键词
                keyword = Keyword(
                    keyword_text=keyword_text,
                    source_topic=source_topic,
                    search_source=search_source,
                    search_round=search_round
                )
                
                session.add(keyword)
                session.commit()
                
                self.logger.debug(f"成功添加关键词: {keyword_text}")
                return True
                
        except IntegrityError:
            return False
        except Exception as e:
            self.logger.error(f"添加关键词失败: {e}")
            return False
    
    def get_unprocessed_keywords(self, limit: int = 100) -> List[str]:
        """
        获取未处理的关键词
        
        Args:
            limit: 获取数量限制
        
        Returns:
            未处理的关键词列表
        """
        try:
            with self.get_session() as session:
                keywords = session.query(Keyword).filter_by(
                    is_processed=False
                ).limit(limit).all()
                
                return [kw.keyword_text for kw in keywords]
                
        except Exception as e:
            self.logger.error(f"获取未处理关键词失败: {e}")
            return []
    
    def mark_keyword_processed(self, keyword_text: str):
        """
        标记关键词为已处理
        
        Args:
            keyword_text: 关键词文本
        """
        try:
            with self.get_session() as session:
                keyword = session.query(Keyword).filter_by(keyword_text=keyword_text).first()
                if keyword:
                    keyword.is_processed = True
                    keyword.updated_at = datetime.now()
                    session.commit()
                    
        except Exception as e:
            self.logger.error(f"标记关键词已处理失败: {e}")
    
    def get_search_history(self, search_query: str) -> Optional[SearchHistory]:
        """
        获取搜索历史
        
        Args:
            search_query: 搜索查询
        
        Returns:
            搜索历史记录或None
        """
        try:
            with self.get_session() as session:
                return session.query(SearchHistory).filter_by(
                    search_query=search_query
                ).first()
                
        except Exception as e:
            self.logger.error(f"获取搜索历史失败: {e}")
            return None
    
    def add_search_history(self, search_query: str, search_round: int, result_count: int):
        """
        添加搜索历史
        
        Args:
            search_query: 搜索查询
            search_round: 搜索轮数
            result_count: 结果数量
        """
        try:
            with self.get_session() as session:
                history = SearchHistory(
                    search_query=search_query,
                    search_round=search_round,
                    result_count=result_count
                )
                
                session.add(history)
                session.commit()
                
        except Exception as e:
            self.logger.error(f"添加搜索历史失败: {e}")
    
    def get_topic_count(self) -> int:
        """获取话题总数"""
        try:
            with self.get_session() as session:
                return session.query(Topic).count()
        except Exception as e:
            self.logger.error(f"获取话题总数失败: {e}")
            return 0
    
    def get_keyword_count(self, processed_only: bool = False) -> int:
        """
        获取关键词总数
        
        Args:
            processed_only: 是否只统计已处理的
        
        Returns:
            关键词总数
        """
        try:
            with self.get_session() as session:
                query = session.query(Keyword)
                if processed_only:
                    query = query.filter_by(is_processed=True)
                return query.count()
        except Exception as e:
            self.logger.error(f"获取关键词总数失败: {e}")
            return 0
    
    def get_redis_lock(self, key: str, timeout: int = 10) -> bool:
        """
        获取Redis分布式锁
        
        Args:
            key: 锁的键名
            timeout: 锁超时时间（秒）
        
        Returns:
            是否获取到锁
        """
        try:
            lock_key = f"lock:{key}"
            result = self.redis_client.set(
                lock_key, 
                threading.get_ident(), 
                ex=timeout, 
                nx=True
            )
            return result is not None
        except Exception as e:
            self.logger.error(f"获取Redis锁失败: {e}")
            return False
    
    def release_redis_lock(self, key: str):
        """
        释放Redis分布式锁
        
        Args:
            key: 锁的键名
        """
        try:
            lock_key = f"lock:{key}"
            self.redis_client.delete(lock_key)
        except Exception as e:
            self.logger.error(f"释放Redis锁失败: {e}")
    
    def close(self):
        """关闭数据库连接"""
        try:
            if hasattr(self, 'engine'):
                self.engine.dispose()
            if hasattr(self, 'redis_pool'):
                self.redis_pool.disconnect()
            self.logger.info("数据库连接已关闭")
        except Exception as e:
            self.logger.error(f"关闭数据库连接失败: {e}")


def create_database_manager(config: Dict) -> DatabaseManager:
    """
    创建数据库管理器实例
    
    Args:
        config: 配置字典
    
    Returns:
        数据库管理器实例
    """
    return DatabaseManager(config)
